from flask import Flask, request, render_template, redirect, url_for, flash
import os
from werkzeug.utils import secure_filename
from datetime import datetime
from predict import predict_image

app = Flask(__name__, template_folder='../templates', static_folder='../static')
app.config['UPLOAD_FOLDER'] = '../uploads'
app.secret_key = '859025462'  # 请替换为你自己的密钥

# 检查并创建上传文件夹
if not os.path.exists(app.config['UPLOAD_FOLDER']):
    os.makedirs(app.config['UPLOAD_FOLDER'])

# 用于存储上传记录的全局变量
upload_history = []


@app.route('/')
def index():
    return render_template('index.html')


@app.route('/upload', methods=['GET', 'POST'])
def upload():
    if request.method == 'POST':
        if 'file' not in request.files:
            flash('没有文件部分')
            return redirect(request.url)
        file = request.files['file']
        if file.filename == '':
            flash('未选择文件')
            return redirect(request.url)
        if file:
            filename = secure_filename(file.filename)
            file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
            file.save(file_path)
            predictions, additional_message = predict_image(file_path)

            # 保存记录 (只记录最可能的结果)
            if predictions:
                top_prediction = predictions[0]
                upload_record = {
                    'timestamp': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                    'filename': filename,
                    'prediction': top_prediction['class_name'],
                    'probability': f"{top_prediction['probability'] * 100:.2f}%",
                    'suggestion': top_prediction['suggestion'],
                    'additional_message': additional_message
                }
                upload_history.append(upload_record)

            return render_template('upload.html', predictions=predictions, additional_message=additional_message)
    return render_template('upload.html')


@app.route('/about')
def about():
    return render_template('about.html')


@app.route('/results')
def results():
    return render_template('results.html', results=upload_history)


if __name__ == '__main__':
    app.run(host='0.0.0.0', debug=True)
